from pydantic import BaseModel
from typing import Dict, List

class Snapshot(BaseModel):
    code: str # 股票代码
    market: int # 市场
    last_price: float # 最新价
    date: int # 日期 yyyyMMdd
    time: int # 分钟 HHmm

class Kline(BaseModel):
    code: str
    market: int
    close: float # 最新价
    open: float # 开盘价
    high: float # 最高价
    low: float # 最低价
    volume: int # 成交量
    date: int # 日期 yyyyMMdd
    time: int # 分钟 HHmm
    close_ema5: float
    close_ema10: float
    amount: float # 成交额
    
class Position(BaseModel):
    code: str
    market: int
    cost_price: float # 成本价
    qty: int # 持仓量
    last_price: float # 现价

class Order(BaseModel):
    order_id: int # 策略内唯一
    code: str
    market: int
    direction: int # 买卖方向，1-买，0-卖
    price: float # 委托价格
    qty: int # 委托数量
    date: int
    time: int
    
class Trade(BaseModel):
    order_id: int 
    code: str
    market: int
    direction: int
    order_price: float # 委托价格
    order_qty: int # 委托数量
    deal_price: float # 成交价格
    deal_qty: int # 成交数量
    date: int
    time: int
    
class Account(BaseModel):
    aid: str
    money: float
    position: Dict[str, Position]
    orders: List[Order]
    trades: List[Trade]
    assets: Dict[int, float]
    date: int
    
class Context:
    def __init__(self, account:Account):
        self.account = account
        
    def has_position(self, code):
        return code in self.account.position
    
    def get_position(self, code):
        return self.account.position[code]
    
    def order(self, order:Order):
        self.account.orders.append(order)
        if order.direction == 1: # 买
            print(f"buy order:{order}")
            need_money = order.price * order.qty
            if need_money < self.account.money:
                self.account.money -= need_money
                if order.code in self.account.position:
                    exist_position = self.account.position[order.code]
                    exist_position.cost_price = (exist_position.cost_price * exist_position.qty + need_money)/(exist_position.qty + order.qty)
                    exist_position.qty += order.qty
                    self.account.position[order.code] = exist_position
                else:
                    new_position = Position(
                        code=order.code,
                        market=order.market,
                        cost_price=order.price,
                        qty=order.qty,
                        last_price=order.price
                    )
                    self.account.position[order.code] = new_position
                return True
            else:
                print("No money")
        else:
            print(f"sell order:{order}")
            if order.code in self.account.position:
                exist_position = self.account.position[order.code]
                if exist_position.qty >= order.qty:
                    in_money = order.qty * order.price
                    self.account.money += in_money
                    exist_position.qty -= order.qty
                    if exist_position.qty <= 0:
                        del self.account.position[order.code]
                    else:
                        self.account.position[order.code] = exist_position
                    return True
                else:
                    print("Not enough qty")
                print(f"money:{self.account.money}")
        return False
        